// Information gathering in Rust .!

use std::fmt::format;
use std::fs::{read_to_string, File, OpenOptions};
use std::io::{self, Write};
use std::path::Path;
use std::process::Command;

use regex::Regex;
use walkdir::WalkDir;
use winreg::enums::{HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE};
use winreg::RegKey;


pub(crate) fn gather_all_info(path: &Path) -> io::Result<()> {
    let mut file = OpenOptions::new().write(true).create(true).open(path)?;

    gather_sys_info(&mut file)?;
    gather_proc_info(&mut file)?;
    gather_user_info(&mut file)?;
    gather_service_info(&mut file)?;
    gather_network_info(&mut file)?;
    gather_filesystem_info(&mut file)?;
    gather_startup_info(&mut file)?;
    gather_disk_info(&mut file)?;
    gather_regis_info(&mut file)?;
    // This may take some time so i have commented it ! . When on Simulation uncomment it based on your needs ! 
    // gather_file_info(&mut file)?;
    save_group_policy_related_checks(&mut file)?;
    Ok(())
}


// Bytes to human readable format conversion !
fn bytes_hr_format(bytes: u64) -> String{
    const KB: u64 = 1024;
    const MB: u64 = KB * 1024;
    const GB: u64 = MB * 1024;
    const TB: u64 = GB * 1024;

    match bytes {
        bytes if bytes >= TB => format!("{:.2} TB", bytes as f64 / TB as f64),
        bytes if bytes >= GB => format!("{:.2} GB", bytes as f64 / GB as f64),
        bytes if bytes >= MB => format!("{:.2} MB", bytes as f64 / MB as f64),
        bytes if bytes >= KB => format!("{:.2} KB", bytes as f64 / KB as f64),
        _ => format!("{} B", bytes), 
    }
}

fn gather_sys_info(file: &mut File) -> io::Result<()>{
    writeln!(file, "######################## System Information ########################")?;
    let output = Command::new("cmd")
        .args(&["/C", "systeminfo"])
        .output()?;
    writeln!(file, "{}", String::from_utf8_lossy(&output.stdout))?;
    Ok(())
}

fn gather_proc_info(file: &mut File) -> io::Result<()>{
    writeln!(file, "######################## Process Information ########################")?;
    let output = Command::new("cmd")
        .args(&["/C", "tasklist"])
        .output()?;
    writeln!(file, "{}", String::from_utf8_lossy(&output.stdout))?;
    Ok(())
}

fn gather_user_info(file: &mut File) -> io::Result<()> {
    writeln!(file, "######################## Users/Groups ########################")?;
    let output = Command::new("cmd")
        .args(&["/C", "net user"])
        .output()?;
    writeln!(file, "{}", String::from_utf8_lossy(&output.stdout))?;
    Ok(())
}

fn gather_service_info(file: &mut File) -> io::Result<()> {
    writeln!(file, "######################## Service Information ########################")?;
    let output = Command::new("cmd")
        .args(&["/C", "sc query"])
        .output()?;
    writeln!(file, "{}", String::from_utf8_lossy(&output.stdout))?;
    Ok(())
}

fn gather_network_info(file: &mut File) -> io::Result<()> {
    writeln!(file, "######################## Network Information ########################")?;
    let output = Command::new("cmd")
        .args(&["/C", "ipconfig /all"])
        .output()?;
    writeln!(file, "{}", String::from_utf8_lossy(&output.stdout))?;
    Ok(())
}

fn gather_filesystem_info(file: &mut File) -> io::Result<()> {
    writeln!(file, "######################## Filesystem Information ########################")?;
    let output = Command::new("cmd")
        .args(&["/C", "fsutil fsinfo drives"])
        .output()?;
    writeln!(file, "{}", String::from_utf8_lossy(&output.stdout))?;
    Ok(())
}

fn gather_startup_info(file: &mut File) -> io::Result<()> {
    writeln!(file, "######################## Startup and Shutdown Information ########################")?;
    let output = Command::new("cmd")
        .args(&["/C", "wmic startup"])
        .output()?;
    writeln!(file, "{}", String::from_utf8_lossy(&output.stdout))?;
    Ok(())
}

fn gather_disk_info(file: &mut File) -> io::Result<()>{
    writeln!(file, "######################## Disk Information ########################")?;
    
    let output = Command::new("cmd")
        .args(&["/C", "wmic logicaldisk get size,freespace,caption"])
        .output()?;

        let output_str = String::from_utf8_lossy(&output.stdout);
        
        let lines: Vec<&str> = output_str.lines().collect();

    if lines.len() > 1{
        for line in &lines[1..] {
            let parts: Vec<&str> = line.split_whitespace().collect();
            if parts.len() == 3{
                let caption = parts[0];
                let free_space: u64 = parts[1].parse().unwrap_or(0);
                let size: u64 = parts[2].parse().unwrap_or(0);

                let hr_format = bytes_hr_format(free_space);
                let hr_size = bytes_hr_format(size);
                
                writeln!(
                    file, "Drive {}: Free Space: {}, TotalSize: {}",
                    caption, hr_format, hr_size
                )?;   
            }
        }
    }
    Ok(())
}

fn gather_regis_info(file: &mut File) -> io::Result<()>{
    writeln!(file, "######################## Registry Password Check ########################")?;

    let hkcu = RegKey::predef(HKEY_CURRENT_USER);
    let hklm  = RegKey::predef(HKEY_LOCAL_MACHINE);

    let reg_path = vec![hkcu, hklm];

    let regex_search = vec![Regex::new(r"(?i)password").unwrap()];

    for reg_key in reg_path {
        for subkey in reg_key.enum_keys().flatten() {
            if let Ok(subkey) = reg_key.open_subkey(&subkey){
                for value in subkey.enum_values().flatten(){
                    let value_name = value.0;
                    let value_data: String = value.1.to_string();
                    
                    for regex in &regex_search{
                        if regex.is_match(&value_name) || regex.is_match(&value_data) {
                            writeln!(file, "Possible Password found in Registry Key: {:?} Value: {}",
                                subkey,
                                value_name,
                            )?;
                        }
                    }
                }
            }
        }
    }
    writeln!(file, "Checked Registry")?;

    Ok(())
}

// This Func will take time so i have uncommented it. Please check and remove the comment for your basics needs !
/* 

fn gather_file_info(file: &mut File) -> io::Result<()>{
    writeln!(file, "######################## File/Folder Check ########################")?;

    // # => For Recursive Search on All Drives !

    // let mut buffer:[u8; 256] = [0; 256];
    // let len = unsafe{
    //     winapi::um::winbase::GetLogicalDriveStringsA(buffer.len() as u32, buffer.as_mut_ptr() as *mut i8)
    // };

    // let mut drives = Vec::new();

    // if len > 0 {
    //     // length is must
    //     let driver_str = String::from_utf8_lossy(&buffer[..len as usize]);
    //     for drive in driver_str.split('\0'){
    //         if !drive.is_empty(){
    //             drives.push(drive.to_string());
    //         }
    //     }
    // }

    // add your own drive letters ! 
    // lets keep this one as testing ! 
    let drives = vec!["C:/"];
    let file_extensions = vec![".xls",  ".xlsx", ".xlsm"];
    let regex_search = vec![Regex::new(r"(?i)password").unwrap(), Regex::new(r"(?i)user").unwrap()];


    for drive in drives{
        // if the condition is ok proceed futher 
        for entry in WalkDir::new(drive).into_iter().filter_map(|e| e.ok()){
            let path = entry.path();
            if path.is_file(){
                let path_str = path.to_string_lossy();

                for ext in &file_extensions{
                    if path_str.ends_with(ext){
                        writeln!(file, "Excel file Found: {}", path_str)?;
                    }
                }

                let content = read_to_string(path);
                if let Ok(content) = content{
                    for regex in &regex_search{
                        if regex.is_match(&content) {
                            writeln!(file, "Possible password found in the File: {}", path_str)?;
                        }
                    }
                }
            }
        }
    }
    Ok(())
}

*/


fn save_group_policy_related_checks(file: &mut File) -> io::Result<()> {
    writeln!(file, "######################## GROUP POLICY RELATED CHECKS ########################\n")?;
    
    writeln!(file, "=========|| SAM / SYSTEM Backup Checks")?;
    let sam_paths = [
        r"C:\Windows\repair\SAM",
        r"C:\Windows\System32\config\RegBack\SAM",
        r"C:\Windows\System32\config\SAM",
        r"C:\Windows\repair\system",
        r"C:\Windows\System32\config\SYSTEM",
        r"C:\Windows\System32\config\RegBack\system",
    ];
    
    for path in &sam_paths {
        if Path::new(path).exists() {
            writeln!(file, "{} Found!", path)?;
        }
    }
    
    writeln!(file, "\n=========|| Group Policy Password Check")?;
    
    let group_policy_files = [
        "Groups.xml", 
        "Services.xml", 
        "Scheduledtasks.xml", 
        "DataSources.xml", 
        "Printers.xml", 
        "Drives.xml"
    ];
    
    let history_paths = [
        r"C:\Windows\System32\GroupPolicy\history",
        r"C:\Documents and Settings\All Users\Application Data\Microsoft\Group Policy\history"
    ];
    
    for history_path in &history_paths {
        if Path::new(history_path).exists() {
            for gp_file in &group_policy_files {
                let output = Command::new("cmd")
                    .args(&["/C", "dir", "/S", &format!("{}\\{}", history_path, gp_file)])
                    .output()?;
                writeln!(file, "{}", String::from_utf8_lossy(&output.stdout))?;
            }
        }
    }
    
    
    Ok(())
}
